package pl.rgolcz.antnet;

import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Random;

import edu.uci.ics.jung.graph.Graph;

public class AntNet
{
    /**
     * Alpha parameter used in pxy probability calculations 
     */
    double alpha = 1.0;
    
    /**
     * Pheromone evaporation coefficient (in the range <0, 1>)
     * 0 - no vaporation at all
     * 1 - pheromone disappears completely each vaporation (useless) 
     */
    double vapour = 0.1;
    
    /**
     * Beta parameter used in pxy probability calculations 
     */
    double beta = 1.0;
    
    /**
     * Q parameter used in pheromone calulations (the higher the value of Q, the more pheromone will ant leave)   
     */
    double Q = 1.0;
    
    /**
     * If memory is on - the ant will choose only from the vertices that were not yet visited
     */
    boolean memory = true;
    
    /**
     * Pheromone matrix - it's bidirectional (separate value for direction a-b and b-a)
     */
    private double[][] pheromone;
    private Graph<Integer, IWeightedEdge> graph;
    private int root;
    private int dest;
    private Random rnd = new Random(new Date().getTime());
    
    private List<Double> solutionLength = new ArrayList<Double>(200);

    private double lastAntLength;
    
    public AntNet(Graph<Integer, IWeightedEdge> graph, int root, int dest)
    {
        this.graph = graph;
        this.root = root;
        this.dest = dest;
    }
    
    /**
     * Single ant simulation
     * @return
     */
    private List<Integer> runAnt()
    {
        double length = 0.0;
        List<Integer> track = new ArrayList<>();
        
        int current = root;
        
        while(current != dest)
        {
            track.add(current);
            
            // Calculate probability for each incident vertex
            List<Integer> incidentVertices = new ArrayList<>();
            List<Double> probability = new ArrayList<>();
            List<Double> distance = new ArrayList<>();
            
            for(IWeightedEdge edge : graph.getIncidentEdges(current))
            {
                int y = graph.getOpposite(current, edge);
                if(memory)
                {
                    if(track.contains(y))
                    {
                        continue;
                    }
                }
                
                double dxy = edge.getWeight();
                double fxy = pheromone[current][y];
                double pxy = Math.pow(1.0 / dxy, alpha) * Math.pow(fxy, beta);
                
                incidentVertices.add(y);
                probability.add(pxy);
                distance.add(dxy);
            }
            
            // Check if there is way out
            if(memory)
            {
                if(incidentVertices.size() == 0)
                {
                    // No way out it's a vertex with single edge
                    //System.out.print(track);
                    //System.out.println(" - NO WAY");
                    return null;                    
                }
            }
            else
            {
                if(incidentVertices.size() == 1)
                {
                    // No way out it's a vertex with single edge
                    //System.out.print(track);
                    //System.out.println(" - NO WAY");
                    return null;
                }
            }
            
            // Convert probability table to incremental probability - this way we can use it for random pick
            // [2 1 4 7] -> [2 3 7 14]
            for(int i = 1; i < probability.size(); i++)
            {
                probability.set(i, probability.get(i) + probability.get(i-1)); 
            }
            
            // Pick the next vertex using provided probability table
            double rand = rnd.nextDouble() * probability.get(probability.size() - 1);
            int i = 0;
            for(; i < probability.size(); i++)
            {
                if(rand <= probability.get(i))
                {
                    break;
                }
            }
            
            // i is the index of the chosen vertex
            int next = incidentVertices.get(i);
            length += distance.get(i);
            
            current = next;
        }
        
        track.add(current);
        
        //System.out.print(track);
        //System.out.println(" - OK");
        
        lastAntLength = length;
        
        return track;
    }
    
    private void init()
    {
        solutionLength.clear();
        
        /* Init pheromone matrix */
        pheromone = new double[graph.getVertexCount()][];
        for(int i = 0; i < pheromone.length; i++)
        {
            pheromone[i] = new double[graph.getVertexCount()];
            for(int k = 0; k < pheromone.length; k++)
            {
                pheromone[i][k] = 1.0;
            }
        }
    }
    
    public void run()
    {
        init();
        
        for(int k = 0; k < 100; k++)
        {
            List<Integer> track = runAnt();
            //double length = 1.0;
            
            /* Pheromone evaporation (each 10 iterations) */
            if(k % 10 == 0)
            {
                for(int i = 0; i < pheromone.length; i++)
                {
                    for(int j = 0; j < pheromone.length; j++)
                    {
                        pheromone[i][j] = Math.max(pheromone[i][j]*(1.0 - vapour), 1.0);
                    }
                }
            }
            
            /* Update pheromone values */
            if(track != null && track.size() >= 2)
            {
                int previous = track.get(0);
                for(int i = 1; i < track.size(); i++)
                {
                    pheromone[previous][track.get(i)] += Q / lastAntLength;
                    previous = track.get(i);
                }
            }
            
            checkCurrentSolution();
        }
    }

    /**
     * Check if any solution exists in current state (the most probable track is current solution).
     * Start in root vertex. Follow the edges with highest probability pxy.
     * If it leads to destination vertex - this is current solution.
     * If not (if there is a loop) - there is no solution.
     * @return lenght of current solution
     */
    private double checkCurrentSolution()
    {
        double length = 0.0;
        List<Integer> track = new ArrayList<>();
        
        int current = root;
        
        while(current != dest)
        {
            track.add(current);
            
            double maxProb = 0.0;
            int nextVert = -1;
            double currentDxy = 0.0;
            
            for(IWeightedEdge edge : graph.getIncidentEdges(current))
            {
                int y = graph.getOpposite(current, edge);
                double dxy = edge.getWeight();
                double fxy = pheromone[current][y];
                double pxy = Math.pow(1.0 / dxy, alpha) * Math.pow(fxy, beta);
                if(pxy > maxProb)
                {
                    maxProb = pxy;
                    nextVert = y;
                    currentDxy = dxy;
                }
            }
            
            length += currentDxy;
            
            /* If there is a loop, solution does not exist */
            if(track.contains(nextVert))
            {
                track.add(nextVert);
                //System.out.println("Solution does not exist - loop: " + track);
                solutionLength.add(null);
                return Double.POSITIVE_INFINITY;
            }
            
            current = nextVert;
        }
        
        track.add(dest);
        
        //System.out.println("Solution exists (" + length + ") " + track);
        solutionLength.add(length);
        return length;
    }
    
    public List<Double> getSolutionLengthList()
    {
        return solutionLength;
    }
    
    public Double getSolutionLength()
    {
        return solutionLength.get(solutionLength.size() - 1);
    }

    public void setAlpha(double alpha)
    {
        this.alpha = alpha;
    }

    public void setVapour(double vapour)
    {
        this.vapour = vapour;
    }

    public void setBeta(double beta)
    {
        this.beta = beta;
    }
    
    public void setQ(double Q)
    {
        this.Q = Q;
    }

    public void setMemory(boolean memory)
    {
        this.memory = memory;
    }
}
